"""
PyG版本的移除自环操作，与dgl.transform.remove_self_loop函数兼容
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from torch_geometric.utils import remove_self_loops as pyg_remove_self_loops
from ops.src.graph.graph import Graph


def remove_self_loops(graph: Graph) -> Graph:
    """
    移除图中的自环，与dgl.transform.remove_self_loop函数兼容
    
    Args:
        graph: 输入图
        
    Returns:
        移除自环后的图
    """
    print("===================remove_self_loops===================")
    # 获取边索引
    edge_index = graph.edge_index
    
    # 使用PyG的remove_self_loops函数移除自环
    edge_index_new, edge_attr_dict = pyg_remove_self_loops(edge_index)
    
    # 创建新图
    new_graph = graph.clone()
    new_graph.edge_index = edge_index_new
    
    # 更新边属性
    if hasattr(graph, 'edge_attr') and graph.edge_attr is not None:
        mask = edge_attr_dict['edge_mask']
        for key in graph.keys:
            if key.startswith('edge_'):
                attr_name = key
                if hasattr(graph, attr_name) and getattr(graph, attr_name) is not None:
                    setattr(new_graph, attr_name, getattr(graph, attr_name)[mask])
    
    return new_graph 